'''
Author: your name
Date: 1970-01-01 08:00:00
LastEditTime: 2020-10-22 16:14:40
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: /Pointnet_Pointnet2/data_utils/myModelNetDataLoader.py
'''

import numpy as np
import warnings
import h5py
from torch.utils.data import Dataset
import torch

warnings.filterwarnings('ignore')

def load_h5(h5_filename):
    f = h5py.File(h5_filename)
    data = f['data'][:]
    label = f['label'][:]
    seg = []
    return (data, label, seg)


def load_data(dir, classification=True):
    data_train0, label_train0, Seglabel_train0 = load_h5(dir + 'ply_data_train0.h5')
    data_train1, label_train1, Seglabel_train1 = load_h5(dir + 'ply_data_train1.h5')
    data_train2, label_train2, Seglabel_train2 = load_h5(dir + 'ply_data_train2.h5')
    data_train3, label_train3, Seglabel_train3 = load_h5(dir + 'ply_data_train3.h5')
    data_train4, label_train4, Seglabel_train4 = load_h5(dir + 'ply_data_train4.h5')
    data_test0, label_test0, Seglabel_test0 = load_h5(dir + 'ply_data_test0.h5')
    data_test1, label_test1, Seglabel_test1 = load_h5(dir + 'ply_data_test1.h5')
    train_data = np.concatenate([data_train0, data_train1, data_train2, data_train3, data_train4])
    train_label = np.concatenate([label_train0, label_train1, label_train2, label_train3, label_train4])
    train_Seglabel = np.concatenate([Seglabel_train0, Seglabel_train1, Seglabel_train2, Seglabel_train3, Seglabel_train4])
    test_data = np.concatenate([data_test0, data_test1])
    test_label = np.concatenate([label_test0, label_test1])
    test_Seglabel = np.concatenate([Seglabel_test0, Seglabel_test1])

    if classification:
        return [train_data, np.squeeze(train_label)], [test_data, np.squeeze(test_label)]
    else:
        return (train_data, np.squeeze(train_Seglabel)), (test_data, np.squeeze(test_Seglabel))


class ModelNetDataLoader(Dataset):
    def __init__(self, dataset, args):
        self.data, self.labels = dataset
        self.args = args

    def __len__(self):
        return len(self.data)


    def __getitem__(self, index):
        return self.data[index], self.labels[index]


class get_loader():
    def __init__(self, args):
        super(get_loader, self).__init__()
        self.data_path = args.data_path
        self.batch_size = args.batch_size
        self.num_worker = args.num_worker
        self.args = args
    
    def create_dataloader(self):
        train_data, test_data = load_data(self.data_path, classification=True)

        TRAIN_DATASET = ModelNetDataLoader(train_data,self.args)
        TEST_DATASET = ModelNetDataLoader(test_data,self.args)

        train_loader = torch.utils.data.DataLoader(
            TRAIN_DATASET, 
            batch_size=self.batch_size, 
            shuffle=True,
            # num_workers=self.num_worker,
            # collate_fn=TRAIN_DATASET.collate_func,
            )
        test_loader = torch.utils.data.DataLoader(
            TEST_DATASET, 
            batch_size=self.batch_size, 
            shuffle=False,
            # num_workers=self.num_worker,
            # collate_fn=TEST_DATASET.collate_func,
            )

        return train_loader,test_loader

